package org.nlp.algo.classifier.svm;

import java.io.File;
import java.io.FileFilter;
import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.TreeMap;

import libsvm.svm;
import libsvm.svm_model;
import libsvm.svm_node;
import libsvm.svm_parameter;
import libsvm.svm_problem;

import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.log4j.Logger;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.util.Version;
import org.nlp.lucene.BloomAnalyzer;

import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;

public class SVMLibClassifier {
	private static Logger _log = Logger.getLogger(SVMLibClassifier.class);
	// 使用一个双向MAP 。
	private HashBiMap<Double, String> types = HashBiMap.create();
	private BiMap<String, Double> retypes = HashBiMap.create();

	//TODO 没有想好已知分类的初始化
	public SVMLibClassifier(File file) throws IOException{
		initType(file);
	}
	
	/**
	 * 需要初始化类型和double
	 * @param file
	 * @throws IOException 
	 */
	private void initType(File file) throws IOException {
		
		List<String> l = FileUtils.readLines(file);
		for (String string : l) {
			types.put(Double.valueOf(string.split(" ")[0]), string.split(" ")[1]);
		}
		retypes = types.inverse();
	}


	/**
	 * 提供训练模型
	 * 
	 * @param TDIR
	 *            训练数据目录 目录结构应该固定 一个目录一个分类。
	 */
	public svm_model train(File trainFile) throws IOException {
		//initType(file);
		svm_parameter param = new svm_parameter();
		initParam(param);
		List<MySVMFile> mySVMFiles = readTransText(trainFile);
		svm_problem prob = new svm_problem();
		// mySvmFile 转制为problem
		trans(prob, mySVMFiles);
		svm_model model = svm.svm_train(prob, param);
		return model;
	}

	public void saveModel(String modelPath, svm_model model) throws IOException {
		svm.svm_save_model(modelPath, model);
	}

	/**
	 * 转换过程，避免了运算过程中要预先知道维度的问题
	 * 
	 * @param prob
	 * @param msfs
	 */
	private void trans(svm_problem prob, List<MySVMFile> msfs) {
		prob.l = msfs.size();
		prob.y = new double[prob.l];
		int max = 0;
		for (int i = 0; i < msfs.size(); i++) {
			max = Math.max(msfs.get(i).getKvs().size(), max);
			prob.y[i] = msfs.get(i).getType();
		}
		prob.x = new svm_node[prob.l][max];
		for (int i = 0; i < msfs.size(); i++) {
			TreeMap<Integer, Double> kvs = msfs.get(i).getKvs();
			Set<Integer> keys = kvs.keySet();
			int j = 0;
			for (Integer key : keys) {
				svm_node x = new svm_node();
				x.index = key;
				x.value = kvs.get(key);
				prob.x[i][j++] = x;
			}
			for (; j < max;) {
				svm_node x = new svm_node();
				x.index = 0;
				x.value = 0;
				prob.x[i][j++] = x;
			}
		}
	}

	
	/**
	 * 分类
	 * @param modelFile
	 * @param text
	 * @return
	 * @throws IOException
	 */
	public String classifier(String modelFile,String text) throws IOException{
		return classifier(svm.svm_load_model(modelFile), text);
	}
	
	/**
	 * 分类
	 * @param model
	 * @param text
	 * @return
	 * @throws IOException
	 */
	public String classifier(svm_model model,String text) throws IOException{
		BloomAnalyzer analyzer = new BloomAnalyzer(Version.LUCENE_40);
		List<svm_node> sns = new ArrayList<svm_node>();
		TokenStream token = analyzer.tokenStream("", new StringReader(
				text));
		CharTermAttribute term = token
				.addAttribute(CharTermAttribute.class);
		while (token.incrementToken()) {
			String word = term.toString();
			// 特征码
			String fe = ReadFeature.features.get(word);
			if (StringUtils.isNotBlank(fe)) {
				svm_node sn = new svm_node();
				sn.index = Integer.parseInt(fe);
				sn.value = 1;
				sns.add(sn);
			}
		}
		analyzer.close();
		Collections.sort(sns);
		svm_node[] x = new svm_node[sns.size()];
		int i =0;
		for (svm_node svm_node : sns) {
			x[i++] = svm_node;
		}
		double d = svm.svm_predict(model, x);
		return types.get(d);
		
		
		
	}

	// TODO KONGXP 理论上我们应该提供set方法 但是没有想好如何封装。
	// 暂时先手工配置参数。
	private static void initParam(svm_parameter param) {
		// default values
		// 多个分类
		param.svm_type = svm_parameter.C_SVC;
		param.kernel_type = svm_parameter.RBF;
		param.degree = 3;
		param.gamma = 0.5;
		param.coef0 = 0;
		param.nu = 0.5;
		param.cache_size = 40;
		param.C = 100;
		// param.C = 1;
		param.eps = 1e-3;
		param.p = 0.1;
		param.shrinking = 1;
		param.probability = 0;
		param.nr_weight = 0;
		param.weight_label = new int[0];
		param.weight = new double[0];
	}

	// 将特征和类型绑定关系
	private List<MySVMFile> readTransText(File file) throws IOException {
		List<MySVMFile> svmFiles = new ArrayList<MySVMFile>();
		File[] types = file.listFiles(new FileFilter() {
			@Override
			public boolean accept(File type) {
				if (type.getName().startsWith(".") || type.isFile()) {
					return false;
				}
				return true;
			}
		});
		BloomAnalyzer analyzer = new BloomAnalyzer(Version.LUCENE_40);
		for (int i = 0; i < types.length; i++) {
			File[] txts = types[i].listFiles(new FileFilter() {
				@Override
				public boolean accept(File pathname) {
					if (pathname.getName().endsWith(".txt")) {
						return true;
					}
					return false;
				}
			});

			// 读取文件，并进行切词然后放入文件中
			for (int j = 0; j < txts.length; j++) {
				_log.debug("read file "+types[i].getName());
				MySVMFile msf = new MySVMFile(retypes.get(types[i].getName()));
				String str = FileUtils.readFileToString(txts[j], "GBK");
				TokenStream token = analyzer.tokenStream("", new StringReader(
						str));
				CharTermAttribute term = token
						.addAttribute(CharTermAttribute.class);
				while (token.incrementToken()) {
					String word = term.toString();
					// 特征码
					String fe = ReadFeature.features.get(word);
					if (StringUtils.isNotBlank(fe)) {
						msf.addPair(Integer.parseInt(fe), 1);
					}
				}
				svmFiles.add(msf);
			}
		}
		analyzer.close();
		return svmFiles;
	}
}
